package edu.buaa.agent;

import edu.buaa.framework.method2.ClassModifier;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

public class JvmAgentWithByteModify {
    public static void premain(String agentArgs, Instrumentation inst) {
        inst.addTransformer(new TimeTransformer(), true);
    }
}

class TimeTransformer implements ClassFileTransformer {
    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
                            ProtectionDomain protectionDomain, byte[] classfileBuffer) {
        byte[] modifyBytes = classfileBuffer;
        if(className.startsWith("edu/buaa/demo"))  {
            // 2. new ClassModifier，并传入需要被修改的字节数组
            ClassModifier cm = new ClassModifier(classfileBuffer);

    // 3. 调用ClassModifier#modifyUTF8Constant修改
            modifyBytes = cm.modifyUTF8Constant("java/lang/System","edu/buaa/framework/PlaybackSystem");
        }
        if ("java/lang/System".equals(className)) {
            try {
                ClassPool cp = ClassPool.getDefault();
                CtClass cc = cp.get("java.lang.System");

                // Modify the currentTimeMillis method
                CtMethod method = cc.getDeclaredMethod("currentTimeMillis");
                method.insertBefore("{ long start = System.nanoTime(); }");
                method.insertAfter("{ System.out.println(\"Execution time: \" + (System.nanoTime() - start) + \" ns\"); }");

                return cc.toBytecode();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return modifyBytes;
    }
}

